{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Seq2Seq.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.8"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "6q7dDydT7n2d",
        "outputId": "be47d969-c754-4cdd-f6d5-db210bcb30b6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 255
        }
      },
      "source": [
        "!wget http://www.manythings.org/anki/cmn-eng.zip\n",
        "!unzip -d ./cmn-eng cmn-eng.zip"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-05-08 02:10:53--  http://www.manythings.org/anki/cmn-eng.zip\n",
            "Resolving www.manythings.org (www.manythings.org)... 104.24.108.196, 104.24.109.196, 2606:4700:3037::6818:6cc4, ...\n",
            "Connecting to www.manythings.org (www.manythings.org)|104.24.108.196|:80... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 963049 (940K) [application/zip]\n",
            "Saving to: ‘cmn-eng.zip’\n",
            "\n",
            "cmn-eng.zip         100%[===================>] 940.48K  2.50MB/s    in 0.4s    \n",
            "\n",
            "2020-05-08 02:10:54 (2.50 MB/s) - ‘cmn-eng.zip’ saved [963049/963049]\n",
            "\n",
            "Archive:  cmn-eng.zip\n",
            "  inflating: ./cmn-eng/cmn.txt       \n",
            "  inflating: ./cmn-eng/_about.txt    \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Zo7QCacp7n2s",
        "colab": {}
      },
      "source": [
        "seed = 2020"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "05LtExs57n20",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n",
        "import time\n",
        "import math\n",
        "import random"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B7h691pQdmyW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QW9iTQeF7n2-"
      },
      "source": [
        "# 读取数据"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "9TBhsBe27n3A",
        "outputId": "3cbf44b8-3dc7-490d-810a-3195ba8f38ad",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "# 每一行数据如下\n",
        "# 'Hi.\\t嗨。\\tCC-BY 2.0 (France) Attribution: tatoeba.org #538123 (CM) & #891077 (Martha)'\n",
        "with open('./cmn-eng/cmn.txt', 'r', encoding='utf-8') as f:\n",
        "    data = f.read()\n",
        "data = data.strip()\n",
        "data = data.split('\\n')\n",
        "print('样本数:\\n', len(data))\n",
        "print('\\n样本示例:')\n",
        "data[0]"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "样本数:\n",
            " 22075\n",
            "\n",
            "样本示例:\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'Hi.\\t嗨。\\tCC-BY 2.0 (France) Attribution: tatoeba.org #538123 (CM) & #891077 (Martha)'"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "vXXhro447n3F",
        "outputId": "c1808dd4-ee5d-4d8a-eed6-90f54bf5256b",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "# 分割英文数据和中文数据\n",
        "en_data = [line.split('\\t')[0] for line in data]\n",
        "ch_data = [line.split('\\t')[1] for line in data]\n",
        "print('英文数据:\\n', en_data[:10])\n",
        "print('\\n中文数据:\\n', ch_data[:10])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "英文数据:\n",
            " ['Hi.', 'Hi.', 'Run.', 'Wait!', 'Wait!', 'Hello!', 'I try.', 'I won!', 'Oh no!', 'Cheers!']\n",
            "\n",
            "中文数据:\n",
            " ['嗨。', '你好。', '你用跑的。', '等等！', '等一下！', '你好。', '让我来。', '我赢了。', '不会吧。', '乾杯!']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "k1ioMRh17n3O",
        "outputId": "bd96dea7-767b-4a21-c1dc-64e6a55b31dd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 102
        }
      },
      "source": [
        "# 按字符级切割，并添加<eos>\n",
        "en_token_list = [[char for char in line]+[\"<eos>\"] for line in en_data]\n",
        "ch_token_list = [[char for char in line]+[\"<eos>\"] for line in ch_data]\n",
        "print('英文数据:\\n', en_token_list[:2])\n",
        "print('\\n中文数据:\\n', ch_token_list[:2])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "英文数据:\n",
            " [['H', 'i', '.', '<eos>'], ['H', 'i', '.', '<eos>']]\n",
            "\n",
            "中文数据:\n",
            " [['嗨', '。', '<eos>'], ['你', '好', '。', '<eos>']]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "wPAejxm17n3T",
        "colab": {}
      },
      "source": [
        "# 基本字典\n",
        "basic_dict = {'<pad>':0, '<unk>':1, '<bos>':2, '<eos>':3}\n",
        "# 分别生成中英文字典 \n",
        "en_vocab = set(''.join(en_data))\n",
        "en2id = {char:i+len(basic_dict) for i, char in enumerate(en_vocab)}\n",
        "en2id.update(basic_dict)\n",
        "id2en = {v:k for k,v in en2id.items()}\n",
        "\n",
        "# 分别生成中英文字典 \n",
        "ch_vocab = set(''.join(ch_data))\n",
        "ch2id = {char:i+len(basic_dict) for i, char in enumerate(ch_vocab)}\n",
        "ch2id.update(basic_dict)\n",
        "id2ch = {v:k for k,v in ch2id.items()}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "jLIjhvZt7n3e",
        "outputId": "96e35b74-b2d5-424e-ac0d-425090c11275",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "source": [
        "# 利用字典，映射数据 \n",
        "en_num_data = [[en2id[en] for en in line ] for line in en_token_list]\n",
        "ch_num_data = [[ch2id[ch] for ch in line] for line in ch_token_list]\n",
        "\n",
        "print('char:', en_data[1])\n",
        "print('index:', en_num_data[1])"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "char: Hi.\n",
            "index: [4, 33, 51, 3]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "IZ5ZtQfN7n3j"
      },
      "source": [
        "# 表示为Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "PPTYSlOw7n3k",
        "colab": {}
      },
      "source": [
        "class TranslationDataset(Dataset):\n",
        "    def __init__(self, src_data, trg_data):\n",
        "        self.src_data = src_data\n",
        "        self.trg_data = trg_data\n",
        "\n",
        "        assert len(src_data) == len(trg_data), \\\n",
        "            \"numbers of src_data  and trg_data must be equal!\"\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.src_data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        src_sample =self.src_data[idx]\n",
        "        src_len = len(self.src_data[idx])\n",
        "        trg_sample = self.trg_data[idx]\n",
        "        trg_len = len(self.trg_data[idx])\n",
        "        return {\"src\": src_sample, \"src_len\": src_len, \"trg\": trg_sample, \"trg_len\": trg_len}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "yUKfKfpc7n3p",
        "colab": {}
      },
      "source": [
        "def padding_batch(batch):\n",
        "    \"\"\"\n",
        "    input: -> list of dict\n",
        "        [{'src': [1, 2, 3], 'trg': [1, 2, 3]}, {'src': [1, 2, 2, 3], 'trg': [1, 2, 2, 3]}]\n",
        "    output: -> dict of tensor \n",
        "        {\n",
        "            \"src\": [[1, 2, 3, 0], [1, 2, 2, 3]].T\n",
        "            \"trg\": [[1, 2, 3, 0], [1, 2, 2, 3]].T\n",
        "        }\n",
        "    \"\"\"\n",
        "    src_lens = [d[\"src_len\"] for d in batch]\n",
        "    trg_lens = [d[\"trg_len\"] for d in batch]\n",
        "    \n",
        "    src_max = max([d[\"src_len\"] for d in batch])\n",
        "    trg_max = max([d[\"trg_len\"] for d in batch])\n",
        "    for d in batch:\n",
        "        d[\"src\"].extend([en2id[\"<pad>\"]]*(src_max-d[\"src_len\"]))\n",
        "        d[\"trg\"].extend([ch2id[\"<pad>\"]]*(trg_max-d[\"trg_len\"]))\n",
        "    srcs = torch.tensor([pair[\"src\"] for pair in batch], dtype=torch.long, device=device)\n",
        "    trgs = torch.tensor([pair[\"trg\"] for pair in batch], dtype=torch.long, device=device)\n",
        "    \n",
        "    batch = {\"src\":srcs.T, \"src_len\":src_lens, \"trg\":trgs.T, \"trg_len\":trg_lens}\n",
        "    return batch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aw8EeZBB9mrK"
      },
      "source": [
        "# 不用Attention机制"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "F19j3Svn7n3v"
      },
      "source": [
        "## 模型"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "8qsKXeDY7n3w",
        "colab": {}
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout=0.5, bidirectional=True):\n",
        "        super(Encoder, self).__init__()\n",
        "        \n",
        "        self.hid_dim = hid_dim\n",
        "        self.n_layers = n_layers\n",
        "        \n",
        "        self.embedding = nn.Embedding(input_dim, emb_dim)\n",
        "        self.gru = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, bidirectional=bidirectional)\n",
        "        \n",
        "    def forward(self, input_seqs, input_lengths, hidden):\n",
        "        # input_seqs = [seq_len, batch]\n",
        "        embedded = self.embedding(input_seqs)\n",
        "        # embedded = [seq_len, batch, embed_dim]\n",
        "        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths, enforce_sorted=False)\n",
        "        \n",
        "        outputs, hidden = self.gru(packed, hidden)        \n",
        "        outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(outputs)\n",
        "        # outputs = [seq_len, batch, hid_dim * n directions]\n",
        "        # output_lengths = [batch]\n",
        "        return outputs, hidden"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3WdZchKu7n33",
        "colab": {}
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout=0.5, bidirectional=True):\n",
        "        super(Decoder, self).__init__()\n",
        "        \n",
        "        self.output_dim = output_dim\n",
        "        self.hid_dim = hid_dim\n",
        "        self.n_layers = n_layers\n",
        "        \n",
        "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
        "        self.gru = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, bidirectional=bidirectional)\n",
        "        \n",
        "        if bidirectional:\n",
        "            self.fc_out = nn.Linear(hid_dim*2, output_dim)\n",
        "        else:\n",
        "            self.fc_out = nn.Linear(hid_dim, output_dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.softmax = nn.LogSoftmax(dim=1)\n",
        "        \n",
        "    def forward(self, token_inputs, hidden):\n",
        "        # token_inputs = [batch]\n",
        "        batch_size = token_inputs.size(0)\n",
        "        embedded = self.dropout(self.embedding(token_inputs).view(1, batch_size, -1))\n",
        "        # embedded = [1, batch, emb_dim]\n",
        "\n",
        "        output, hidden = self.gru(embedded, hidden)\n",
        "        # output = [1, batch,  n_directions * hid_dim]\n",
        "        # hidden = [n_layers * n_directions, batch, hid_dim]\n",
        "        \n",
        "        output = self.fc_out(output.squeeze(0))\n",
        "        output = self.softmax(output)\n",
        "        # output = [batch, output_dim]\n",
        "        return output, hidden"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "tkDLZQVR7n4A",
        "colab": {}
      },
      "source": [
        "class Seq2Seq(nn.Module):\n",
        "    def __init__(self, \n",
        "                 encoder, \n",
        "                 decoder, \n",
        "                 device, \n",
        "                 predict=False, \n",
        "                 basic_dict=None,\n",
        "                 max_len=100\n",
        "                 ):\n",
        "        super(Seq2Seq, self).__init__()\n",
        "        \n",
        "        self.device = device\n",
        "\n",
        "        self.encoder = encoder\n",
        "        self.decoder = decoder\n",
        "\n",
        "        self.predict = predict  # 训练阶段还是预测阶段\n",
        "        self.basic_dict = basic_dict  # decoder的字典，存放特殊token对应的id\n",
        "        self.max_len = max_len  # 翻译时最大输出长度\n",
        "\n",
        "        self.enc_n_layers = self.encoder.gru.num_layers\n",
        "        self.enc_n_directions = 2 if self.encoder.gru.bidirectional else 1\n",
        "        self.dec_n_directions = 2 if self.decoder.gru.bidirectional else 1\n",
        "\n",
        "        assert encoder.hid_dim == decoder.hid_dim, \\\n",
        "            \"Hidden dimensions of encoder and decoder must be equal!\"\n",
        "        assert encoder.n_layers == decoder.n_layers, \\\n",
        "            \"Encoder and decoder must have equal number of layers!\"\n",
        "        assert self.enc_n_directions >= self.dec_n_directions, \\\n",
        "            \"If decoder is bidirectional, encoder must be bidirectional either!\"\n",
        "        \n",
        "    def forward(self, input_batches, input_lengths, target_batches=None, target_lengths=None, teacher_forcing_ratio=0.5):\n",
        "        # input_batches = target_batches = [seq_len, batch]\n",
        "        batch_size = input_batches.size(1)\n",
        "        \n",
        "        BOS_token = self.basic_dict[\"<bos>\"]\n",
        "        EOS_token = self.basic_dict[\"<eos>\"]\n",
        "        PAD_token = self.basic_dict[\"<pad>\"]\n",
        "\n",
        "        # 初始化\n",
        "        encoder_hidden = torch.zeros(self.enc_n_layers*self.enc_n_directions, batch_size, self.encoder.hid_dim, device=self.device)\n",
        "        \n",
        "        # encoder_output = [seq_len, batch, hid_dim * n directions]\n",
        "        # encoder_hidden = [n_layers*n_directions, batch, hid_dim]\n",
        "        encoder_output, encoder_hidden = self.encoder(\n",
        "            input_batches, input_lengths, encoder_hidden)\n",
        "\n",
        "        # 初始化\n",
        "        decoder_input = torch.tensor([BOS_token] * batch_size, dtype=torch.long, device=self.device)\n",
        "        if self.enc_n_directions == self.dec_n_directions:\n",
        "            decoder_hidden = encoder_hidden\n",
        "        else:\n",
        "            L = encoder_hidden.size(0)\n",
        "            decoder_hidden = encoder_hidden[range(0, L, 2)] + encoder_hidden[range(1, L, 2)]\n",
        "\n",
        "        if self.predict:\n",
        "            # 预测阶段使用\n",
        "            # 一次只输入一句话\n",
        "            assert batch_size == 1, \"batch_size of predict phase must be 1!\"\n",
        "            output_tokens = []\n",
        "\n",
        "            while True:\n",
        "                decoder_output, decoder_hidden = self.decoder(\n",
        "                    decoder_input, decoder_hidden\n",
        "                )\n",
        "                # [1, 1]\n",
        "                topv, topi = decoder_output.topk(1)\n",
        "                decoder_input = topi.squeeze(1)  # 上一个预测作为下一个输入\n",
        "                output_token = topi.squeeze().detach().item()\n",
        "                if output_token == EOS_token or len(output_tokens) == self.max_len:\n",
        "                    break\n",
        "                output_tokens.append(output_token)\n",
        "            return output_tokens\n",
        "\n",
        "        else:\n",
        "            # 训练阶段\n",
        "            max_target_length = max(target_lengths)\n",
        "            all_decoder_outputs = torch.zeros((max_target_length, batch_size, self.decoder.output_dim), device=self.device)\n",
        "\n",
        "            for t in range(max_target_length):\n",
        "                use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
        "                if use_teacher_forcing:\n",
        "                    # decoder_output = [batch, output_dim]\n",
        "                    # decoder_hidden = [n_layers*n_directions, batch, hid_dim]\n",
        "                    decoder_output, decoder_hidden = self.decoder(\n",
        "                        decoder_input, decoder_hidden\n",
        "                    )\n",
        "                    all_decoder_outputs[t] = decoder_output\n",
        "                    decoder_input = target_batches[t]  # 下一个输入来自训练数据\n",
        "                else:\n",
        "                    decoder_output, decoder_hidden = self.decoder(\n",
        "                        decoder_input, decoder_hidden\n",
        "                    )\n",
        "                    # [batch, 1]\n",
        "                    topv, topi = decoder_output.topk(1)\n",
        "                    all_decoder_outputs[t] = decoder_output\n",
        "                    decoder_input = topi.squeeze(1)  # 下一个输入来自模型预测\n",
        "            \n",
        "            loss_fn = nn.NLLLoss(ignore_index=PAD_token)\n",
        "            loss = loss_fn(\n",
        "                all_decoder_outputs.reshape(-1,self.decoder.output_dim ),  # [batch*seq_len, output_dim]\n",
        "                target_batches.reshape(-1)                                                 # [batch*seq_len]\n",
        "            )\n",
        "            return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DGIK_4IO7n4G"
      },
      "source": [
        "## 训练和预测代码"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "TvUaEp0Y7n4H",
        "colab": {}
      },
      "source": [
        "def epoch_time(start_time, end_time):\n",
        "    elapsed_time = end_time - start_time\n",
        "    elapsed_mins = int(elapsed_time / 60)\n",
        "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
        "    return elapsed_mins, elapsed_secs"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "FcB6O0cJ7n4P",
        "colab": {}
      },
      "source": [
        "def train(\n",
        "    model,\n",
        "    data_loader, \n",
        "    optimizer, \n",
        "    clip=1, \n",
        "    teacher_forcing_ratio=0.5, \n",
        "    print_every=None  # None不打印\n",
        "    ):\n",
        "    model.predict = False\n",
        "    model.train()\n",
        "\n",
        "    if print_every == 0:\n",
        "        print_every = 1\n",
        "\n",
        "    print_loss_total = 0  # 每次打印都重置\n",
        "    start = time.time()\n",
        "    epoch_loss = 0\n",
        "    for i, batch in enumerate(data_loader):\n",
        "\n",
        "        # shape = [seq_len, batch]\n",
        "        input_batchs = batch[\"src\"]\n",
        "        target_batchs = batch[\"trg\"]\n",
        "        # list\n",
        "        input_lens = batch[\"src_len\"]\n",
        "        target_lens = batch[\"trg_len\"]\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        loss = model(input_batchs, input_lens, target_batchs, target_lens, teacher_forcing_ratio)\n",
        "        print_loss_total += loss.item()\n",
        "        epoch_loss += loss.item()\n",
        "        loss.backward()\n",
        "\n",
        "        # 梯度裁剪\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
        "\n",
        "        optimizer.step()\n",
        "\n",
        "        if print_every and (i+1) % print_every == 0:\n",
        "            print_loss_avg = print_loss_total / print_every\n",
        "            print_loss_total = 0\n",
        "            print('\\tCurrent Loss: %.4f' % print_loss_avg)\n",
        "\n",
        "    return epoch_loss / len(data_loader)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "DgYv7rsp7n4V",
        "colab": {}
      },
      "source": [
        "def evaluate(\n",
        "    model,\n",
        "    data_loader, \n",
        "    print_every=None\n",
        "    ):\n",
        "    model.predict = False\n",
        "    model.eval()\n",
        "    if print_every == 0:\n",
        "        print_every = 1\n",
        "\n",
        "    print_loss_total = 0  # 每次打印都重置\n",
        "    start = time.time()\n",
        "    epoch_loss = 0\n",
        "    with torch.no_grad():\n",
        "        for i, batch in enumerate(data_loader):\n",
        "\n",
        "            # shape = [seq_len, batch]\n",
        "            input_batchs = batch[\"src\"]\n",
        "            target_batchs = batch[\"trg\"]\n",
        "            # list\n",
        "            input_lens = batch[\"src_len\"]\n",
        "            target_lens = batch[\"trg_len\"]\n",
        "\n",
        "            loss = model(input_batchs, input_lens, target_batchs, target_lens, teacher_forcing_ratio=0)\n",
        "            print_loss_total += loss.item()\n",
        "            epoch_loss += loss.item()\n",
        "\n",
        "            if print_every and (i+1) % print_every == 0:\n",
        "                print_loss_avg = print_loss_total / print_every\n",
        "                print_loss_total = 0\n",
        "                print('\\tCurrent Loss: %.4f' % print_loss_avg)\n",
        "\n",
        "    return epoch_loss / len(data_loader)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ldcZM35r7n4h",
        "colab": {}
      },
      "source": [
        "def translate(\n",
        "    model,\n",
        "    sample, \n",
        "    idx2token=None\n",
        "    ):\n",
        "    model.predict = True\n",
        "    model.eval()\n",
        "\n",
        "    # shape = [seq_len, 1]\n",
        "    input_batch = sample[\"src\"]\n",
        "    # list\n",
        "    input_len = sample[\"src_len\"]\n",
        "\n",
        "    output_tokens = model(input_batch, input_len)\n",
        "    output_tokens = [idx2token[t] for t in output_tokens]\n",
        "\n",
        "    return \"\".join(output_tokens)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "r9luHxDc7n4q"
      },
      "source": [
        "## 开始训练"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "QSLPWwzT7n4s",
        "colab": {}
      },
      "source": [
        "INPUT_DIM = len(en2id)\n",
        "OUTPUT_DIM = len(ch2id)\n",
        "\n",
        "# 超参数\n",
        "BATCH_SIZE = 32\n",
        "ENC_EMB_DIM = 256\n",
        "DEC_EMB_DIM = 256\n",
        "HID_DIM = 512\n",
        "N_LAYERS = 2\n",
        "ENC_DROPOUT = 0.5\n",
        "DEC_DROPOUT = 0.5\n",
        "LEARNING_RATE = 1e-4\n",
        "N_EPOCHS = 200\n",
        "CLIP = 1\n",
        "\n",
        "bidirectional = True\n",
        "enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT, bidirectional)\n",
        "dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT, bidirectional)\n",
        "model = Seq2Seq(enc, dec, device, basic_dict=basic_dict).to(device)\n",
        "\n",
        "## encoder和encoder设置相同的学习策略\n",
        "optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
        "# ## encoder和encoder设置不同的学习策略\n",
        "# optimizer_grouped_parameters = [\n",
        "#         {'params': [p for n, p in model.named_parameters() if 'encoder' in n], 'lr': LEARNING_RATE},\n",
        "#         {'params': [p for n, p in model.named_parameters() if 'decoder' in n], 'lr': LEARNING_RATE*2}\n",
        "# ]\n",
        "# optimizer = optim.Adam(optimizer_grouped_parameters)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "xHDdqB2t7n43",
        "colab": {}
      },
      "source": [
        "# 数据集\n",
        "train_set = TranslationDataset(en_num_data, ch_num_data)\n",
        "train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, collate_fn=padding_batch)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39nVqNwVdm2R",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(N_EPOCHS):\n",
        "    \n",
        "    start_time = time.time()\n",
        "    train_loss = train(model, train_loader, optimizer, CLIP)\n",
        "    valid_loss = evaluate(model, train_loader)\n",
        "    end_time = time.time()\n",
        "    \n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        torch.save(model.state_dict(), 'en2ch-model.pt')\n",
        "\n",
        "    if epoch %2 == 0:\n",
        "        epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
        "        print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
        "        print(f'\\tTrain Loss: {train_loss:.3f} | Val. Loss: {valid_loss:.3f}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7AKhdoGdm2g",
        "colab_type": "text"
      },
      "source": [
        "## 模型使用"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oqXQO8ezdm2h",
        "colab_type": "code",
        "colab": {},
        "outputId": "31464c6f-82ed-46a5-a710-7e01e6f487e0"
      },
      "source": [
        "print(\"best valid loss：\", best_valid_loss)\n",
        "# 加载最优权重\n",
        "model.load_state_dict(torch.load('en2ch-model.pt'))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "best valid loss： 0.13634515966649802\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 111
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "twr9NGHzdm2w",
        "colab_type": "code",
        "colab": {},
        "outputId": "fa4ed3e9-37bf-4f72-c133-a9d93a4c926e"
      },
      "source": [
        "random.seed(seed)\n",
        "for i in random.sample(range(len(en_num_data)), 10):  # 随机看10个\n",
        "    en_tokens = list(filter(lambda x: x!=0, en_num_data[i]))  # 过滤零\n",
        "    ch_tokens = list(filter(lambda x: x!=3 and x!=0, ch_num_data[i]))  # 和机器翻译作对照\n",
        "    sentence = [id2en[t] for t in en_tokens]\n",
        "    print(\"【原文】\")\n",
        "    print(\"\".join(sentence))\n",
        "    translation = [id2ch[t] for t in ch_tokens]\n",
        "    print(\"【原文】\")\n",
        "    print(\"\".join(translation))\n",
        "    test_sample = {}\n",
        "    test_sample[\"src\"] = torch.tensor(en_tokens, dtype=torch.long, device=device).reshape(-1, 1)\n",
        "    test_sample[\"src_len\"] = [len(en_tokens)]\n",
        "    print(\"【机器翻译】\")\n",
        "    print(translate(model, test_sample, id2ch), end=\"\\n\\n\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "【原文】\n",
            "Do you know if Tom has spoken to Mary recently?<eos>\n",
            "【原文】\n",
            "你是否知道最近湯姆跟瑪麗有沒有說過話？\n",
            "【机器翻译】\n",
            "你是否知道最近湯姆跟瑪麗沒有沒有說過話？\n",
            "\n",
            "【原文】\n",
            "Children of six and above should attend school.<eos>\n",
            "【原文】\n",
            "六岁及以上的小孩应该上学。\n",
            "【机器翻译】\n",
            "六岁及以上的小孩应该上学。\n",
            "\n",
            "【原文】\n",
            "She doesn't like soccer.<eos>\n",
            "【原文】\n",
            "她不喜欢足球。\n",
            "【机器翻译】\n",
            "她不喜欢足球。\n",
            "\n",
            "【原文】\n",
            "I'd like to know the phone number of the nearest American Express office.<eos>\n",
            "【原文】\n",
            "我想知道最近的美國運通辦事處的電話號碼。\n",
            "【机器翻译】\n",
            "我想知道最近的美國運通辦處處。\n",
            "\n",
            "【原文】\n",
            "I've never seen him wearing jeans.<eos>\n",
            "【原文】\n",
            "我從來沒有看過他穿牛仔褲。\n",
            "【机器翻译】\n",
            "我從來沒有看過他穿牛仔褲。\n",
            "\n",
            "【原文】\n",
            "We had a good time playing chess.<eos>\n",
            "【原文】\n",
            "我們下棋玩得很開心。\n",
            "【机器翻译】\n",
            "我們下棋玩得很開心。\n",
            "\n",
            "【原文】\n",
            "You don't need to work on Sundays.<eos>\n",
            "【原文】\n",
            "星期天的時候，你不用工作。\n",
            "【机器翻译】\n",
            "星期天的時候，你不用工作。\n",
            "\n",
            "【原文】\n",
            "She will be back within a week.<eos>\n",
            "【原文】\n",
            "她一週之內會回來。\n",
            "【机器翻译】\n",
            "她一週之內會回來。\n",
            "\n",
            "【原文】\n",
            "He told me that he was busy then.<eos>\n",
            "【原文】\n",
            "他告诉我他那时很忙。\n",
            "【机器翻译】\n",
            "他告诉我他那时很忙。\n",
            "\n",
            "【原文】\n",
            "Sentences begin with a capital letter.<eos>\n",
            "【原文】\n",
            "句子以一個大寫字母開頭。\n",
            "【机器翻译】\n",
            "句子以一個大寫字母開頭頭。\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ANbudxH17n5Z",
        "outputId": "131e7f2c-2d6b-4deb-9406-6f1368ec0554",
        "colab": {}
      },
      "source": [
        "random.seed(seed*2)\n",
        "for i in random.sample(range(len(en_num_data)), 10):  # 随机看10个\n",
        "    en_tokens = list(filter(lambda x: x!=0, en_num_data[i]))  # 过滤零\n",
        "    ch_tokens = list(filter(lambda x: x!=3 and x!=0, ch_num_data[i]))  # 和机器翻译作对照\n",
        "    sentence = [id2en[t] for t in en_tokens]\n",
        "    print(\"【原文】\")\n",
        "    print(\"\".join(sentence))\n",
        "    translation = [id2ch[t] for t in ch_tokens]\n",
        "    print(\"【原文】\")\n",
        "    print(\"\".join(translation))\n",
        "    test_sample = {}\n",
        "    test_sample[\"src\"] = torch.tensor(en_tokens, dtype=torch.long, device=device).reshape(-1, 1)\n",
        "    test_sample[\"src_len\"] = [len(en_tokens)]\n",
        "    print(\"【机器翻译】\")\n",
        "    print(translate(model, test_sample, id2ch), end=\"\\n\\n\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "【原文】\n",
            "They're in danger.<eos>\n",
            "【原文】\n",
            "他們有危險。\n",
            "【机器翻译】\n",
            "他們有危險。\n",
            "\n",
            "【原文】\n",
            "Somebody is playing the piano.<eos>\n",
            "【原文】\n",
            "有人正在彈鋼琴。\n",
            "【机器翻译】\n",
            "有人正在彈鋼琴。\n",
            "\n",
            "【原文】\n",
            "The explosion may have been caused by a gas leak.<eos>\n",
            "【原文】\n",
            "這場爆炸有可能是瓦斯外洩引起的。\n",
            "【机器翻译】\n",
            "這場爆炸有可能是瓦斯外洩引起的。\n",
            "\n",
            "【原文】\n",
            "We'll ask Tom and see what he thinks.<eos>\n",
            "【原文】\n",
            "我們要問問湯姆，看看他怎麼想。\n",
            "【机器翻译】\n",
            "我們要問問湯姆，看看他怎麼想。\n",
            "\n",
            "【原文】\n",
            "I'd like a nonstop flight to New York.<eos>\n",
            "【原文】\n",
            "我想要一个直飞纽约的航班。\n",
            "【机器翻译】\n",
            "我想要一个直飞纽约的航班。\n",
            "\n",
            "【原文】\n",
            "I'm afraid the job I've got for you won't be easy.<eos>\n",
            "【原文】\n",
            "我恐怕我給你找的工作不輕鬆。\n",
            "【机器翻译】\n",
            "我恐怕我給你找的工作不輕鬆。\n",
            "\n",
            "【原文】\n",
            "He couldn't stand the bitterness of the coffee.<eos>\n",
            "【原文】\n",
            "他受不了咖啡的苦味。\n",
            "【机器翻译】\n",
            "他受不了咖啡的苦味。\n",
            "\n",
            "【原文】\n",
            "Even if I knew the answer, I wouldn't tell you.<eos>\n",
            "【原文】\n",
            "即使我知道答案，我也不会告诉你。\n",
            "【机器翻译】\n",
            "即使我知道答案，我也不会告诉你。\n",
            "\n",
            "【原文】\n",
            "He threw a stone into the pond.<eos>\n",
            "【原文】\n",
            "他扔一塊石頭到池塘裡。\n",
            "【机器翻译】\n",
            "他扔一塊石頭到池塘裡。\n",
            "\n",
            "【原文】\n",
            "The fuel tank in the car is full.<eos>\n",
            "【原文】\n",
            "汽車油箱是滿的。\n",
            "【机器翻译】\n",
            "汽車油箱是滿的。\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Es1g11r09g9v"
      },
      "source": [
        "# Attention机制"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0h3IhpdS7n5j"
      },
      "source": [
        "## 模型"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "SBpm6azx_FLe",
        "colab": {}
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout=0.5, bidirectional=True):\n",
        "        super(Encoder, self).__init__()\n",
        "        \n",
        "        self.hid_dim = hid_dim\n",
        "        self.n_layers = n_layers\n",
        "        \n",
        "        self.embedding = nn.Embedding(input_dim, emb_dim)\n",
        "        self.gru = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, bidirectional=bidirectional)\n",
        "        \n",
        "    def forward(self, input_seqs, input_lengths, hidden):\n",
        "        # input_seqs = [seq_len, batch]\n",
        "        embedded = self.embedding(input_seqs)\n",
        "        # embedded = [seq_len, batch, embed_dim]\n",
        "        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths, enforce_sorted=False)\n",
        "        \n",
        "        outputs, hidden = self.gru(packed, hidden)        \n",
        "        outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(outputs)\n",
        "        # outputs = [seq_len, batch, hid_dim * n directions]\n",
        "        # output_lengths = [batch]\n",
        "        return outputs, hidden"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "rMbsHQa38Kv1",
        "colab": {}
      },
      "source": [
        "class Attn(nn.Module):\n",
        "    def __init__(self, method, hidden_size):\n",
        "        super(Attn, self).__init__()\n",
        "        self.method = method\n",
        "        if self.method not in ['dot', 'general', 'concat']:\n",
        "            raise ValueError(self.method, \"is not an appropriate attention method.\")\n",
        "        self.hidden_size = hidden_size\n",
        "        if self.method == 'general':\n",
        "            self.attn = nn.Linear(self.hidden_size, hidden_size)\n",
        "        elif self.method == 'concat':\n",
        "            self.attn = nn.Linear(self.hidden_size * 2, hidden_size)\n",
        "            self.v = nn.Parameter(torch.FloatTensor(hidden_size))\n",
        "\n",
        "    def dot_score(self, hidden, encoder_output):\n",
        "        return torch.sum(hidden * encoder_output, dim=2)  # [seq_len, batch]\n",
        "\n",
        "    def general_score(self, hidden, encoder_output):\n",
        "        energy = self.attn(encoder_output)  # [seq_len, batch, hid_dim]\n",
        "        return torch.sum(hidden * energy, dim=2)  # [seq_len, batch]\n",
        "\n",
        "    def concat_score(self, hidden, encoder_output):\n",
        "        # hidden.expand(encoder_output.size(0), -1, -1) -> [seq_len, batch, N]\n",
        "        energy = self.attn(torch.cat((hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh()\n",
        "        # energy = [sql_len, batch, hidden_size]\n",
        "        return torch.sum(self.v * energy, dim=2)  # [seq_len, batch]\n",
        "\n",
        "    def forward(self, hidden, encoder_outputs):\n",
        "        # hidden = [1, batch,  n_directions * hid_dim]\n",
        "        # encoder_outputs = [seq_len, batch, hid dim * n directions]\n",
        "        if self.method == 'general':\n",
        "            attn_energies = self.general_score(hidden, encoder_outputs)\n",
        "        elif self.method == 'concat':\n",
        "            attn_energies = self.concat_score(hidden, encoder_outputs)\n",
        "        elif self.method == 'dot':\n",
        "            attn_energies = self.dot_score(hidden, encoder_outputs)\n",
        "\n",
        "        attn_energies = attn_energies.t()  # [batch, seq_len]\n",
        " \n",
        "        return F.softmax(attn_energies, dim=1).unsqueeze(1)  # softmax归一化# [batch, 1, seq_len]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lLbqgxNg8XT4",
        "colab": {}
      },
      "source": [
        "class AttnDecoder(nn.Module):\n",
        "    def __init__(self, output_dim, emb_dim, hid_dim, n_layers=1, dropout=0.5, bidirectional=True, attn_method=\"general\"):\n",
        "        super(AttnDecoder, self).__init__()\n",
        "\n",
        "        self.output_dim = output_dim\n",
        "        self.emb_dim = emb_dim\n",
        "        self.hid_dim = hid_dim\n",
        "        self.n_layers = n_layers\n",
        "        self.dropout = dropout\n",
        "\n",
        "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
        "        self.embedding_dropout = nn.Dropout(dropout)\n",
        "        self.gru = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, bidirectional=bidirectional)\n",
        "        \n",
        "        if bidirectional:\n",
        "            self.concat = nn.Linear(hid_dim * 2 * 2, hid_dim*2)\n",
        "            self.out = nn.Linear(hid_dim*2, output_dim)\n",
        "            self.attn = Attn(attn_method, hid_dim*2)\n",
        "        else:\n",
        "            self.concat = nn.Linear(hid_dim * 2, hid_dim)\n",
        "            self.out = nn.Linear(hid_dim, output_dim)\n",
        "            self.attn = Attn(attn_method, hid_dim)\n",
        "        self.softmax = nn.LogSoftmax(dim=1)\n",
        "\n",
        "    def forward(self, token_inputs, last_hidden, encoder_outputs):\n",
        "        batch_size = token_inputs.size(0)\n",
        "        embedded = self.embedding(token_inputs)\n",
        "        embedded = self.embedding_dropout(embedded)\n",
        "        embedded = embedded.view(1, batch_size, -1) # [1, B, hid_dim]\n",
        "\n",
        "        gru_output, hidden = self.gru(embedded, last_hidden)\n",
        "        # gru_output = [1, batch,  n_directions * hid_dim]\n",
        "        # hidden = [n_layers * n_directions, batch, hid_dim]\n",
        "\n",
        "        # encoder_outputs = [sql_len, batch, hid dim * n directions]\n",
        "        attn_weights = self.attn(gru_output, encoder_outputs)\n",
        "        # attn_weights = [batch, 1, sql_len]\n",
        "        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))\n",
        "        # [batch, 1, hid_dim * n directions]\n",
        "\n",
        "        # LuongAttention\n",
        "        gru_output = gru_output.squeeze(0) # [batch, n_directions * hid_dim]\n",
        "        context = context.squeeze(1)       # [batch, n_directions * hid_dim]\n",
        "        concat_input = torch.cat((gru_output, context), 1)  # [batch, n_directions * hid_dim * 2]\n",
        "        concat_output = torch.tanh(self.concat(concat_input))  # [batch, n_directions*hid_dim]\n",
        "\n",
        "        output = self.out(concat_output)  # [batch, output_dim]\n",
        "        output = self.softmax(output)\n",
        "\n",
        "        return output, hidden, attn_weights"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "y94lk1qX_Cxn",
        "colab": {}
      },
      "source": [
        "class Seq2Seq(nn.Module):\n",
        "    def __init__(self, \n",
        "                 encoder, \n",
        "                 decoder, \n",
        "                 device, \n",
        "                 predict=False, \n",
        "                 basic_dict=None,\n",
        "                 max_len=100\n",
        "                 ):\n",
        "        super(Seq2Seq, self).__init__()\n",
        "        \n",
        "        self.device = device\n",
        "\n",
        "        self.encoder = encoder\n",
        "        self.decoder = decoder\n",
        "\n",
        "        self.predict = predict  # 训练阶段还是预测阶段\n",
        "        self.basic_dict = basic_dict  # decoder的字典，存放特殊token对应的id\n",
        "        self.max_len = max_len  # 翻译时最大输出长度\n",
        "\n",
        "        assert encoder.hid_dim == decoder.hid_dim, \\\n",
        "            \"Hidden dimensions of encoder and decoder must be equal!\"\n",
        "        assert encoder.n_layers == decoder.n_layers, \\\n",
        "            \"Encoder and decoder must have equal number of layers!\"\n",
        "        assert encoder.gru.bidirectional == decoder.gru.bidirectional, \\\n",
        "            \"Decoder and encoder must had same value of bidirectional attribute!\"\n",
        "        \n",
        "    def forward(self, input_batches, input_lengths, target_batches=None, target_lengths=None, teacher_forcing_ratio=0.5):\n",
        "        # input_batches = [seq_len, batch]\n",
        "        # target_batches = [seq_len, batch]\n",
        "        batch_size = input_batches.size(1)\n",
        "        \n",
        "        BOS_token = self.basic_dict[\"<bos>\"]\n",
        "        EOS_token = self.basic_dict[\"<eos>\"]\n",
        "        PAD_token = self.basic_dict[\"<pad>\"]\n",
        "\n",
        "        # 初始化\n",
        "        enc_n_layers = self.encoder.gru.num_layers\n",
        "        enc_n_directions = 2 if self.encoder.gru.bidirectional else 1\n",
        "        encoder_hidden = torch.zeros(enc_n_layers*enc_n_directions, batch_size, self.encoder.hid_dim, device=self.device)\n",
        "        \n",
        "        # encoder_outputs = [input_lengths, batch, hid_dim * n directions]\n",
        "        # encoder_hidden = [n_layers*n_directions, batch, hid_dim]\n",
        "        encoder_outputs, encoder_hidden = self.encoder(\n",
        "            input_batches, input_lengths, encoder_hidden)\n",
        "\n",
        "        # 初始化\n",
        "        decoder_input = torch.tensor([BOS_token] * batch_size, dtype=torch.long, device=self.device)\n",
        "        decoder_hidden = encoder_hidden\n",
        "\n",
        "        if self.predict:\n",
        "            # 一次只输入一句话\n",
        "            assert batch_size == 1, \"batch_size of predict phase must be 1!\"\n",
        "            output_tokens = []\n",
        "\n",
        "            while True:\n",
        "                decoder_output, decoder_hidden, decoder_attn = self.decoder(\n",
        "                    decoder_input, decoder_hidden, encoder_outputs\n",
        "                )\n",
        "                # [1, 1]\n",
        "                topv, topi = decoder_output.topk(1)\n",
        "                decoder_input = topi.squeeze(1).detach()\n",
        "                output_token = topi.squeeze().detach().item()\n",
        "                if output_token == EOS_token or len(output_tokens) == self.max_len:\n",
        "                    break\n",
        "                output_tokens.append(output_token)\n",
        "            return output_tokens\n",
        "\n",
        "        else:\n",
        "            max_target_length = max(target_lengths)\n",
        "            all_decoder_outputs = torch.zeros((max_target_length, batch_size, self.decoder.output_dim), device=self.device)\n",
        "\n",
        "            for t in range(max_target_length):\n",
        "                use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
        "                if use_teacher_forcing:\n",
        "                    # decoder_output = [batch, output_dim]\n",
        "                    # decoder_hidden = [n_layers*n_directions, batch, hid_dim]\n",
        "                    decoder_output, decoder_hidden, decoder_attn = self.decoder(\n",
        "                        decoder_input, decoder_hidden, encoder_outputs\n",
        "                    )\n",
        "                    all_decoder_outputs[t] = decoder_output\n",
        "                    decoder_input = target_batches[t]  # 下一个输入来自训练数据\n",
        "                else:\n",
        "                    decoder_output, decoder_hidden, decoder_attn = self.decoder(\n",
        "                        decoder_input, decoder_hidden, encoder_outputs\n",
        "                    )\n",
        "                    # [batch, 1]\n",
        "                    topv, topi = decoder_output.topk(1)\n",
        "                    all_decoder_outputs[t] = decoder_output\n",
        "                    decoder_input = topi.squeeze(1).detach()  # 下一个输入来自模型预测\n",
        "            \n",
        "            loss_fn = nn.NLLLoss(ignore_index=PAD_token)\n",
        "            loss = loss_fn(\n",
        "                all_decoder_outputs.reshape(-1, self.decoder.output_dim),  # [batch*seq_len, output_dim]\n",
        "                target_batches.reshape(-1)               # [batch*seq_len]\n",
        "            )\n",
        "            return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kLS0uKiu9WtW"
      },
      "source": [
        "## 开始训练"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3ALxD62K8a-s",
        "colab": {}
      },
      "source": [
        "INPUT_DIM = len(en2id)\n",
        "OUTPUT_DIM = len(ch2id)\n",
        "# 超参数\n",
        "BATCH_SIZE = 32\n",
        "ENC_EMB_DIM = 256\n",
        "DEC_EMB_DIM = 256\n",
        "HID_DIM = 512\n",
        "N_LAYERS = 2\n",
        "ENC_DROPOUT = 0.5\n",
        "DEC_DROPOUT = 0.5\n",
        "LEARNING_RATE = 1e-4\n",
        "N_EPOCHS = 200\n",
        "CLIP = 1\n",
        "\n",
        "bidirectional = True\n",
        "attn_method = \"general\"\n",
        "enc = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT, bidirectional)\n",
        "dec = AttnDecoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT, bidirectional, attn_method)\n",
        "model = Seq2Seq(enc, dec, device, basic_dict=basic_dict).to(device)\n",
        "\n",
        "optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "A1HXlUlz9_oV",
        "colab": {}
      },
      "source": [
        "# 数据集\n",
        "train_set = TranslationDataset(en_num_data, ch_num_data)\n",
        "train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, collate_fn=padding_batch)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "GbnYsq9p-GbI",
        "colab": {}
      },
      "source": [
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(N_EPOCHS):\n",
        "    \n",
        "    start_time = time.time()\n",
        "    train_loss = train(model, train_loader, optimizer, CLIP)\n",
        "    valid_loss = evaluate(model, train_loader)\n",
        "    end_time = time.time()\n",
        "    \n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        torch.save(model.state_dict(), \"en2ch-attn-model.pt\")\n",
        "\n",
        "    if epoch %2 == 0:\n",
        "        epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
        "        print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
        "        print(f'\\tTrain Loss: {train_loss:.3f} | Val. Loss: {valid_loss:.3f}')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2aw5nHpOdm5J",
        "colab_type": "code",
        "colab": {},
        "outputId": "f41efc81-8de2-4c84-a4f4-437fda6b5dd6"
      },
      "source": [
        "print(\"best valid loss：\", best_valid_loss)\n",
        "# 加载最优权重\n",
        "model.load_state_dict(torch.load(\"en2ch-attn-model.pt\"))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "best valid loss： 0.12658343204901135\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 21
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8hGVhGBydm5R",
        "colab_type": "code",
        "colab": {},
        "outputId": "0ebf3daa-4717-42aa-f0ed-c535bb5d2248"
      },
      "source": [
        "random.seed(seed)\n",
        "for i in random.sample(range(len(en_num_data)), 10):  # 随机看10个\n",
        "    en_tokens = list(filter(lambda x: x!=0, en_num_data[i]))  # 过滤零\n",
        "    ch_tokens = list(filter(lambda x: x!=3 and x!=0, ch_num_data[i]))  # 和机器翻译作对照\n",
        "    sentence = [id2en[t] for t in en_tokens]\n",
        "    print(\"【原文】\")\n",
        "    print(\"\".join(sentence))\n",
        "    translation = [id2ch[t] for t in ch_tokens]\n",
        "    print(\"【原文】\")\n",
        "    print(\"\".join(translation))\n",
        "    test_sample = {}\n",
        "    test_sample[\"src\"] = torch.tensor(en_tokens, dtype=torch.long, device=device).reshape(-1, 1)\n",
        "    test_sample[\"src_len\"] = [len(en_tokens)]\n",
        "    print(\"【机器翻译】\")\n",
        "    print(translate(model, test_sample, id2ch), end=\"\\n\\n\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "【原文】\n",
            "Do you know if Tom has spoken to Mary recently?<eos>\n",
            "【原文】\n",
            "你是否知道最近湯姆跟瑪麗有沒有說過話？\n",
            "【机器翻译】\n",
            "你是否知道最近湯姆跟瑪麗有沒有說過話？\n",
            "\n",
            "【原文】\n",
            "Children of six and above should attend school.<eos>\n",
            "【原文】\n",
            "六岁及以上的小孩应该上学。\n",
            "【机器翻译】\n",
            "六岁及以上的小孩应该上学。\n",
            "\n",
            "【原文】\n",
            "She doesn't like soccer.<eos>\n",
            "【原文】\n",
            "她不喜欢足球。\n",
            "【机器翻译】\n",
            "她不喜欢足球。\n",
            "\n",
            "【原文】\n",
            "I'd like to know the phone number of the nearest American Express office.<eos>\n",
            "【原文】\n",
            "我想知道最近的美國運通辦事處的電話號碼。\n",
            "【机器翻译】\n",
            "我想知道最近的最通運事處的電話號碼。\n",
            "\n",
            "【原文】\n",
            "I've never seen him wearing jeans.<eos>\n",
            "【原文】\n",
            "我從來沒有看過他穿牛仔褲。\n",
            "【机器翻译】\n",
            "我從來沒有看過他穿牛仔褲。\n",
            "\n",
            "【原文】\n",
            "We had a good time playing chess.<eos>\n",
            "【原文】\n",
            "我們下棋玩得很開心。\n",
            "【机器翻译】\n",
            "我們下棋玩得很開心。\n",
            "\n",
            "【原文】\n",
            "You don't need to work on Sundays.<eos>\n",
            "【原文】\n",
            "星期天的時候，你不用工作。\n",
            "【机器翻译】\n",
            "星期天的時候，你不用工作。\n",
            "\n",
            "【原文】\n",
            "She will be back within a week.<eos>\n",
            "【原文】\n",
            "她一週之內會回來。\n",
            "【机器翻译】\n",
            "她一週之內會回來。\n",
            "\n",
            "【原文】\n",
            "He told me that he was busy then.<eos>\n",
            "【原文】\n",
            "他告诉我他那时很忙。\n",
            "【机器翻译】\n",
            "他告诉我他那时很忙。\n",
            "\n",
            "【原文】\n",
            "Sentences begin with a capital letter.<eos>\n",
            "【原文】\n",
            "句子以一個大寫字母開頭。\n",
            "【机器翻译】\n",
            "句子以一個大寫字母開頭。\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k0_5yNr-dm5a",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}